import pydot
from tree import Tree
from node import Node

class DisplayTree:

    def __init__(self):
        return

    def getTree(self, tree, name):
        graph = pydot.Dot(graph_type='graph')
        node = tree.getRoot().getChild1()
        graph = self.generateGraph(graph, node)
        graph.write_png('C:\Users\Charles\Documents\jira-reporting\Thesis\Images\%s.png' % name)

    def generateGraph(self, graph, node):
        if node.getChild1() is not None:
            graph.add_edge(pydot.Edge(node.getLabel(), node.getChild1().getLabel()))
            self.generateGraph(graph, node.getChild1())
        if node.getChild2() is not None:
            graph.add_edge(pydot.Edge(node.getLabel(), node.getChild2().getLabel()))
            self.generateGraph(graph, node.getChild2())
        return graph
